# visualization/plotter.py
from datetime import datetime

import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
from typing import Dict, Any, List, Optional, Tuple

# 尝试设置支持中文的字体
try:
    # 优先尝试常见的黑体、宋体等
    matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'Songti SC', 'Microsoft YaHei', 'WenQuanYi Micro Hei'] # Unix/Linux/macOS/Windows
    matplotlib.rcParams['axes.unicode_minus'] = False  # 解决负号显示为方块的问题
except Exception as e:
    print(f"设置matplotlib中文字体时出错: {e}。中文可能无法正确显示。")

class Plotter:
    def __init__(self, optimization_results_df: Optional[pd.DataFrame] = None):
        self.results_df = optimization_results_df

    def plot_equity_curve(self, equity_curve_data: List[Tuple[Any, float]], title: str = "权益曲线"):
        if not equity_curve_data:
            print("无权益曲线数据可绘制。")
            return

        # 过滤掉可能的公元1年或其他非常早的初始占位时间戳
        # 假设这些早期时间戳的权益值与初始余额相同
        valid_equity_data = []
        initial_balance_for_filtering = None
        if len(equity_curve_data) > 0:
            # 假设第一个点可能是虚拟的初始点，后续的点才是真实的回测数据
            # 如果第一个点的时间非常早 (例如小于1970年)，并且其权益等于初始余额
            # 我们可能想要跳过它，除非它是唯一的点
            first_dt_obj = None
            if isinstance(equity_curve_data[0][0], str):
                try:
                    first_dt_obj = pd.to_datetime(equity_curve_data[0][0]).to_pydatetime()
                except:
                    pass  # 无法解析则不处理
            elif isinstance(equity_curve_data[0][0], datetime):
                first_dt_obj = equity_curve_data[0][0]

            # 这里的 initial_balance_for_filtering 需要从Account对象获取，或者通过参数传入
            # 为了简化，我们假设如果只有一个点，或者第一个点时间正常，就不过滤
            # 更好的方法是 Account 在记录第一个真实数据点时，替换掉那个虚拟的公元1年的点

            # 简单过滤：只取1900年之后的数据
            min_sensible_year = 1900
            for item_time, item_equity in equity_curve_data:
                dt_to_check = None
                if isinstance(item_time, str):
                    try:
                        dt_to_check = pd.to_datetime(item_time).to_pydatetime()
                    except:
                        continue  # 跳过无法解析的时间
                elif isinstance(item_time, datetime):
                    dt_to_check = item_time
                elif isinstance(item_time, pd.Timestamp):
                    dt_to_check = item_time.to_pydatetime()
                else:
                    continue  # 跳过未知时间类型

                if dt_to_check and dt_to_check.year >= min_sensible_year:
                    valid_equity_data.append((item_time, item_equity))

            if not valid_equity_data and equity_curve_data:  # 如果过滤后为空但原先有数据，则使用原数据
                print("警告：过滤早期时间戳后数据为空，将使用原始权益曲线数据。")
                valid_equity_data = equity_curve_data
            elif not valid_equity_data and not equity_curve_data:
                print("无有效权益曲线数据可绘制。")
                return

        try:
            times = [pd.to_datetime(item[0]) if not isinstance(item[0], pd.Timestamp) else item[0] for item in
                     valid_equity_data]
            equities = [item[1] for item in valid_equity_data]
        except pd.errors.OutOfBoundsDatetime as oob_dt_error:
            print(f"转换权益曲线时间数据时仍发生 OutOfBoundsDatetime 错误: {oob_dt_error}")
            print("将尝试使用原始时间数据直接绘制（可能格式不佳）。")
            times = [item[0] for item in valid_equity_data]
            equities = [item[1] for item in valid_equity_data]
        except Exception as e:
            print(f"转换权益曲线时间数据时出错: {e}. 将尝试直接绘制。")
            times = [item[0] for item in valid_equity_data]
            equities = [item[1] for item in valid_equity_data]

        if not times:
            print("过滤或转换后无时间数据可绘制。")
            return

        plt.figure(figsize=(12, 6))
        plt.plot(times, equities, label="账户权益")
        plt.title(title)
        plt.xlabel("时间")
        plt.ylabel("权益")
        plt.legend()
        plt.grid(True)
        plt.xticks(rotation=45)
        plt.tight_layout()
        plt.show()

    def plot_parameter_plane(
            self,
            param_x_name: str,  # DataFrame中代表x轴参数的列名 (例如 "param_ema_period_5m")
            param_y_name: str,  # DataFrame中代表y轴参数的列名 (例如 "param_rsi_period_5m")
            metric_to_plot_name: str,  # 要在平面上显示的热力图指标 (例如 "total_net_profit" 或 "score")
            title: Optional[str] = None
    ):
        """
        绘制参数平面热力图。
        需要 GridSearchOptimizer.get_results_dataframe() 的输出来准备数据。
        """
        if self.results_df is None or self.results_df.empty:
            print("无优化结果数据可用于绘制参数平面图。")
            return
        if param_x_name not in self.results_df.columns or \
                param_y_name not in self.results_df.columns or \
                metric_to_plot_name not in self.results_df.columns:
            print(f"错误: 提供的参数名或指标名在结果DataFrame中不存在。可用列: {self.results_df.columns.tolist()}")
            return

        try:
            # 创建透视表作为热力图的数据源
            heatmap_data = self.results_df.pivot_table(
                index=param_y_name,  # Y轴
                columns=param_x_name,  # X轴
                values=metric_to_plot_name  # 热力值
            )

            plt.figure(figsize=(10, 8))
            sns.heatmap(heatmap_data, annot=True, fmt=".2f", cmap="viridis")  # annot=True 显示数值, fmt=".2f" 格式化
            plot_title = title if title else f"参数平面: {metric_to_plot_name} vs ({param_x_name}, {param_y_name})"
            plt.title(plot_title)
            plt.xlabel(param_x_name.replace("param_", ""))  # 清理X轴标签
            plt.ylabel(param_y_name.replace("param_", ""))  # 清理Y轴标签
            plt.tight_layout()
            plt.show()
        except Exception as e:
            print(f"绘制参数平面图时出错: {e}")
            import traceback
            traceback.print_exc()

    def plot_3d_surface(
            self,
            param_x_name: str,
            param_y_name: str,
            metric_to_plot_name: str,
            title: Optional[str] = None
    ):
        """绘制参数平面3D曲面图 (需要 matplotlib 的 3D 支持)"""
        if self.results_df is None or self.results_df.empty:
            print("无优化结果数据可用于绘制3D曲面图。")
            return
        # ... (类似的热力图数据准备和检查) ...

        try:
            from mpl_toolkits.mplot3d import Axes3D  # 尝试导入3D绘图工具
            import numpy as np

            pivot_data = self.results_df.pivot_table(
                index=param_y_name, columns=param_x_name, values=metric_to_plot_name
            )

            X = pivot_data.columns.values.astype(float)  # X轴参数值
            Y = pivot_data.index.values.astype(float)  # Y轴参数值
            X, Y = np.meshgrid(X, Y)  # 创建网格
            Z = pivot_data.values  # Z轴指标值

            fig = plt.figure(figsize=(12, 8))
            ax = fig.add_subplot(111, projection='3d')

            ax.plot_surface(X, Y, Z, cmap='viridis', edgecolor='none')

            plot_title = title if title else f"3D参数曲面: {metric_to_plot_name} vs ({param_x_name}, {param_y_name})"
            ax.set_title(plot_title)
            ax.set_xlabel(param_x_name.replace("param_", ""))
            ax.set_ylabel(param_y_name.replace("param_", ""))
            ax.set_zlabel(metric_to_plot_name)
            plt.tight_layout()
            plt.show()
        except ImportError:
            print("绘制3D曲面图需要 mpl_toolkits.mplot3d。")
        except Exception as e:
            print(f"绘制3D曲面图时出错: {e}")
            import traceback
            traceback.print_exc()